from fastai.core import *
from fastai.vision.data import ImageImageList, ImageDataBunch, imagenet_stats
from fastai.vision.transform import get_transforms


def get_colorize_data(
        sz: int,
        bs: int,
        crappy_path: Path,
        good_path: Path,
        random_seed: int = None,
        keep_pct: float = 1.0,
        num_workers: int = 8,
        xtra_tfms=[],
) -> ImageDataBunch:
    src = (
        # 用于图像到图像的任务 加载图片、划分训练集和测试集
        # 工厂方法之一实例化它
        ImageImageList.from_folder(crappy_path, convert_mode='RGB')
            # 仅使用完整数据集的sample_pct样本和可选seed
            .use_partial_data(sample_pct=keep_pct, seed=random_seed)
            # 通过将valid_pct放入验证集中来随机拆分项目，可以传递可选seed
            .split_by_rand_pct(0.1, seed=random_seed)
    )

    data = (
        # 将func应用于每个输入以获取其标签
        src.label_from_func(lambda x: good_path / x.relative_to(crappy_path))
            # 数据增强
            # max_zoom ：如果不等于1或更小，则以概率p_affine应用1.和max_zoom之间的随机缩放
            # max_lighting ：如果不是None，则以max_lighting的概率p_lighting施加由max_lighting控制的随机闪电和对比度变化
            # max_warp ：如果不为None，则以概率p_affine施加-max_warp和maw_warp之间的随机对称幅度幅度
            # xtra_tfms ：您想要应用的其他转换的列表
            .transform(
            get_transforms(
                max_zoom=1.2, max_lighting=0.5, max_warp=0.25, xtra_tfms=xtra_tfms
            ),
            size=sz,
            tfm_y=True,
        )
            .databunch(bs=bs, num_workers=num_workers, no_check=True)
            .normalize(imagenet_stats, do_y=True)
    )

    data.c = 3
    return data


def get_dummy_databunch() -> ImageDataBunch:
    path = Path('./dummy/')
    return get_colorize_data(
        sz=1, bs=1, crappy_path=path, good_path=path, keep_pct=0.001
    )
